#!/usr/bin/env python
# Copyright 2025 Background Incandescent
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# 最先导入基本模块和设置日志
import sys
import os
import logging
import traceback

# 创建日志目录（如果不存在）
os.makedirs("logs", exist_ok=True)

# 设置最基本的日志系统，在导入其他模块前
try:
    logging.basicConfig(
        level=logging.DEBUG,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler("logs/seqkit_mcp_debug.log", mode="w"),
            logging.StreamHandler()
        ]
    )
    logger = logging.getLogger("seqkit-mcp")
    logger.info("日志系统初始化成功")
except Exception as e:
    # 确保即使日志设置失败也能看到错误
    with open("seqkit_startup_error.txt", "w") as f:
        f.write(f"日志初始化失败: {str(e)}\n")
        f.write(traceback.format_exc())
    print(f"日志初始化失败: {str(e)}")
    print(traceback.format_exc())
    sys.exit(1)

try:
    # 导入其他模块
    logger.info("开始导入其他模块")
    import asyncio
    import json
    import subprocess
    import tempfile
    from pathlib import Path
    from typing import Any, List, Optional, Dict
    from pydantic import AnyUrl

    logger.info("导入MCP模块")
    import mcp.types as types
    from mcp.server import Server, NotificationOptions
    from mcp.server.stdio import stdio_server
    from mcp.server.models import InitializationOptions
    logger.info("所有模块导入成功")
except ImportError as e:
    logger.critical(f"导入模块失败: {str(e)}")
    logger.critical(traceback.format_exc())
    print(f"导入模块失败: {str(e)}")
    print("请确保已经安装了所需的包: pip install mcp")
    sys.exit(1)
except Exception as e:
    logger.critical(f"初始化过程中出现意外错误: {str(e)}")
    logger.critical(traceback.format_exc())
    print(f"初始化过程中出现意外错误: {str(e)}")
    sys.exit(1)

# Seqkit MCP工具的提示模板
SEQKIT_PROMPT_TEMPLATE = """
<mcp>
这个服务器提供了Seqkit生物序列分析工具的访问能力。Seqkit是一个用于FASTA/Q文件处理的强大工具集。

工具:
1. "统计序列信息" (seqkit://stats) - 分析序列文件获取统计信息，如序列数量、长度、GC含量等
   参数: input_file (必需) - 输入文件路径

2. "获取序列" (seqkit://seq) - 对序列进行操作，如反向、互补或移除间隙
   参数: 
   - input_file (必需) - 输入文件路径
   - reverse (可选) - 是否反向序列
   - complement (可选) - 是否生成互补序列
   - remove_gaps (可选) - 是否移除间隙
   - output_file (可选) - 输出文件路径

3. "查找序列" (seqkit://grep) - 根据模式查找序列
   参数:
   - input_file (必需) - 输入文件路径
   - pattern (必需) - 搜索模式
   - by_name (可选) - 是否在序列名称中搜索
   - by_seq (可选) - 是否在序列内容中搜索
   - invert_match (可选) - 是否反向匹配
   - output_file (可选) - 输出文件路径

4. "转换序列格式" (seqkit://convert) - 在不同序列格式之间转换
   参数:
   - input_file (必需) - 输入文件路径
   - output_format (必需) - 输出格式，可选：fasta, fastq
   - full_header (可选) - 是否保留完整标题
   - output_file (可选) - 输出文件路径

5. "翻译核酸序列" (seqkit://translate) - 将核酸序列翻译为蛋白质序列
   参数:
   - input_file (必需) - 输入文件路径
   - frame (可选) - 翻译框架
   - trim (可选) - 是否修剪序列
   - clean (可选) - 是否清理序列
   - output_file (可选) - 输出文件路径

这些工具帮助用户方便地分析和操作生物序列数据，无需编写复杂的命令行。
</mcp>

你可以通过这个MCP服务器使用Seqkit工具分析生物序列数据。例如，你可以：

1. 获取序列文件的基本统计信息
2. 执行序列操作，如反向、互补
3. 根据模式查找特定序列
4. 转换序列格式
5. 将DNA/RNA序列翻译为蛋白质序列

要使用这些功能，你可以直接告诉我你想做什么，我会调用适当的工具。

例如，你可以说"请分析example.fasta文件中的序列统计信息"，或者"请将example.fasta中的序列反向互补并保存到output.fasta"。
"""

# 确保seqkit可执行
try:
    SEQKIT_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "seqkit")
    if not os.path.exists(SEQKIT_PATH) and sys.platform == "win32":
        SEQKIT_PATH += ".exe"

    logger.info(f"Seqkit路径: {SEQKIT_PATH}")
    logger.info(f"文件存在: {os.path.exists(SEQKIT_PATH)}")
    
    # 测试seqkit是否可用
    if os.path.exists(SEQKIT_PATH):
        try:
            version_result = subprocess.run(
                [SEQKIT_PATH, "version"],
                capture_output=True,
                text=True,
                check=True
            )
            logger.info(f"Seqkit版本: {version_result.stdout.strip()}")
        except subprocess.CalledProcessError as e:
            logger.warning(f"无法获取Seqkit版本: {e.stderr}")
    else:
        logger.warning(f"Seqkit可执行文件不存在: {SEQKIT_PATH}")
except Exception as e:
    logger.error(f"检查Seqkit时出错: {str(e)}")
    logger.error(traceback.format_exc())

try:
    # 创建MCP服务器
    logger.info("初始化MCP服务器")
    app = Server("seqkit", version="1.0.0")
    logger.info("MCP服务器初始化成功")

    # 定义资源列表
    @app.list_resources()
    async def list_resources() -> List[types.Resource]:
        """列出可用的Seqkit资源和工具"""
        logger.info("列出资源")
        try:
            return [
                types.Resource(
                    uri=AnyUrl("seqkit://stats"),
                    name="统计序列信息",
                    description="分析序列文件获取统计信息，如序列数量、长度、GC含量等",
                    mimeType="text/plain"
                ),
                types.Resource(
                    uri=AnyUrl("seqkit://seq"),
                    name="获取序列",
                    description="对序列进行操作，如反向、互补或移除间隙",
                    mimeType="text/plain"
                ),
                types.Resource(
                    uri=AnyUrl("seqkit://grep"),
                    name="查找序列",
                    description="根据模式查找序列",
                    mimeType="text/plain"
                ),
                types.Resource(
                    uri=AnyUrl("seqkit://convert"),
                    name="转换序列格式",
                    description="在不同序列格式之间转换",
                    mimeType="text/plain"
                ),
                types.Resource(
                    uri=AnyUrl("seqkit://translate"),
                    name="翻译核酸序列",
                    description="将核酸序列翻译为蛋白质序列",
                    mimeType="text/plain"
                )
            ]
        except Exception as e:
            logger.error(f"列出资源时出错: {str(e)}")
            logger.error(traceback.format_exc())
            return []

    # 定义提示模板
    @app.list_prompts()
    async def list_prompts() -> List[types.Prompt]:
        """列出可用的提示模板"""
        logger.info("列出提示模板")
        return [
            types.Prompt(
                name="seqkit-intro",
                description="关于如何使用Seqkit MCP服务器的介绍",
                arguments=[]
            )
        ]

    @app.get_prompt()
    async def get_prompt(name: str, arguments: Dict[str, str] | None) -> types.GetPromptResult:
        """获取提示模板内容"""
        logger.info(f"获取提示模板: {name}")
        if name != "seqkit-intro":
            raise ValueError(f"未知的提示模板: {name}")

        return types.GetPromptResult(
            description="Seqkit MCP服务器介绍",
            messages=[
                types.PromptMessage(
                    role="user",
                    content=types.TextContent(type="text", text=SEQKIT_PROMPT_TEMPLATE.strip())
                )
            ]
        )

    # 定义工具列表
    @app.list_tools()
    async def list_tools() -> List[types.Tool]:
        """列出所有可用工具"""
        logger.info("列出工具")
        return [
            types.Tool(
                name="统计序列信息",
                description="分析序列文件获取统计信息，如序列数量、长度、GC含量等",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "input_file": {"type": "string", "description": "输入文件路径"}
                    },
                    "required": ["input_file"]
                }
            ),
            types.Tool(
                name="获取序列",
                description="对序列进行操作，如反向、互补或移除间隙",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "input_file": {"type": "string", "description": "输入文件路径"},
                        "reverse": {"type": "boolean", "description": "是否反向序列"},
                        "complement": {"type": "boolean", "description": "是否生成互补序列"},
                        "remove_gaps": {"type": "boolean", "description": "是否移除间隙"},
                        "output_file": {"type": "string", "description": "输出文件路径"}
                    },
                    "required": ["input_file"]
                }
            ),
            types.Tool(
                name="查找序列",
                description="根据模式查找序列",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "input_file": {"type": "string", "description": "输入文件路径"},
                        "pattern": {"type": "string", "description": "搜索模式"},
                        "by_name": {"type": "boolean", "description": "是否在序列名称中搜索"},
                        "by_seq": {"type": "boolean", "description": "是否在序列内容中搜索"},
                        "invert_match": {"type": "boolean", "description": "是否反向匹配"},
                        "output_file": {"type": "string", "description": "输出文件路径"}
                    },
                    "required": ["input_file", "pattern"]
                }
            ),
            types.Tool(
                name="转换序列格式",
                description="在不同序列格式之间转换",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "input_file": {"type": "string", "description": "输入文件路径"},
                        "output_format": {"type": "string", "description": "输出格式，可选：fasta, fastq"},
                        "full_header": {"type": "boolean", "description": "是否保留完整标题"},
                        "output_file": {"type": "string", "description": "输出文件路径"}
                    },
                    "required": ["input_file", "output_format"]
                }
            ),
            types.Tool(
                name="翻译核酸序列",
                description="将核酸序列翻译为蛋白质序列",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "input_file": {"type": "string", "description": "输入文件路径"},
                        "frame": {"type": "integer", "description": "翻译框架"},
                        "trim": {"type": "boolean", "description": "是否修剪序列"},
                        "clean": {"type": "boolean", "description": "是否清理序列"},
                        "output_file": {"type": "string", "description": "输出文件路径"}
                    },
                    "required": ["input_file"]
                }
            )
        ]

    # 统一的工具调用处理函数
    @app.call_tool()
    async def call_tool(name: str, arguments: Dict[str, Any] | None) -> List[types.TextContent]:
        """处理工具调用请求"""
        logger.info(f"调用工具: {name}, 参数: {arguments}")
        
        if not arguments:
            return [types.TextContent(type="text", text="错误: 未提供参数")]
            
        try:
            if name == "统计序列信息":
                result = await handle_stats(arguments)
                return [types.TextContent(type="text", text=str(result))]
                
            elif name == "获取序列":
                result = await handle_seq(arguments)
                return [types.TextContent(type="text", text=str(result))]
                
            elif name == "查找序列":
                result = await handle_grep(arguments)
                return [types.TextContent(type="text", text=str(result))]
                
            elif name == "转换序列格式":
                result = await handle_convert(arguments)
                return [types.TextContent(type="text", text=str(result))]
                
            elif name == "翻译核酸序列":
                result = await handle_translate(arguments)
                return [types.TextContent(type="text", text=str(result))]
                
            else:
                return [types.TextContent(type="text", text=f"错误: 未知工具 {name}")]
        except Exception as e:
            logger.error(f"工具调用错误: {str(e)}")
            logger.error(traceback.format_exc())
            return [types.TextContent(type="text", text=f"错误: {str(e)}")]

    # 工具实现函数
    async def handle_stats(params: Dict[str, Any]) -> Dict[str, Any]:
        """处理统计序列请求"""
        logger.info(f"处理stats请求: {params}")
        
        if "input_file" not in params:
            return {"error": "需要提供输入文件路径(input_file)"}
        
        input_file = params["input_file"]
        
        try:
            result = subprocess.run(
                [SEQKIT_PATH, "stats", input_file],
                capture_output=True,
                text=True,
                check=True
            )
            return {"result": result.stdout}
        except subprocess.CalledProcessError as e:
            logger.error(f"执行stats命令时出错: {str(e)}")
            return {"error": f"执行错误: {e.stderr if e.stderr else str(e)}"}
        except Exception as e:
            logger.error(f"处理stats请求时出错: {str(e)}")
            logger.error(traceback.format_exc())
            return {"error": f"处理请求时出错: {str(e)}"}
    
    async def handle_seq(params: Dict[str, Any]) -> Dict[str, Any]:
        """处理序列操作请求"""
        logger.info(f"处理seq请求: {params}")
        
        if "input_file" not in params:
            return {"error": "需要提供输入文件路径(input_file)"}
        
        input_file = params["input_file"]
        args = [SEQKIT_PATH, "seq"]
        
        if params.get("reverse"):
            args.append("-r")
        if params.get("complement"):
            args.append("-p") 
        if params.get("remove_gaps"):
            args.append("-g")
        
        output_file = params.get("output_file")
        if output_file:
            args.extend(["-o", output_file])
        
        try:
            if output_file:
                result = subprocess.run(
                    args + [input_file],
                    capture_output=True,
                    text=True,
                    check=True
                )
                return {"result": f"已保存到 {output_file}", "message": result.stdout}
            else:
                with tempfile.NamedTemporaryFile(mode="w+", delete=False) as temp:
                    temp_name = temp.name
                
                subprocess.run(
                    args + ["-o", temp_name, input_file],
                    check=True
                )
                
                with open(temp_name, "r") as f:
                    content = f.read()
                
                os.unlink(temp_name)
                return {"result": content}
        except subprocess.CalledProcessError as e:
            logger.error(f"执行seq命令时出错: {str(e)}")
            return {"error": f"执行错误: {e.stderr if e.stderr else str(e)}"}
        except Exception as e:
            logger.error(f"处理seq请求时出错: {str(e)}")
            logger.error(traceback.format_exc())
            return {"error": f"处理请求时出错: {str(e)}"}
    
    async def handle_grep(params: Dict[str, Any]) -> Dict[str, Any]:
        """处理序列搜索请求"""
        logger.info(f"处理grep请求: {params}")
        
        if "input_file" not in params:
            return {"error": "需要提供输入文件路径(input_file)"}
        if "pattern" not in params:
            return {"error": "需要提供搜索模式(pattern)"}
        
        input_file = params["input_file"]
        pattern = params["pattern"]
        
        args = [SEQKIT_PATH, "grep"]
        
        args.extend(["-p", pattern])
        
        if params.get("by_name"):
            args.append("-n")
        if params.get("by_seq"):
            args.append("-s")
        if params.get("invert_match"):
            args.append("-v")
        
        output_file = params.get("output_file")
        if output_file:
            args.extend(["-o", output_file])
        
        try:
            if output_file:
                result = subprocess.run(
                    args + [input_file],
                    capture_output=True,
                    text=True,
                    check=True
                )
                return {"result": f"已保存到 {output_file}", "message": result.stdout}
            else:
                with tempfile.NamedTemporaryFile(mode="w+", delete=False) as temp:
                    temp_name = temp.name
                
                subprocess.run(
                    args + ["-o", temp_name, input_file],
                    check=True
                )
                
                with open(temp_name, "r") as f:
                    content = f.read()
                
                os.unlink(temp_name)
                return {"result": content}
        except subprocess.CalledProcessError as e:
            logger.error(f"执行grep命令时出错: {str(e)}")
            return {"error": f"执行错误: {e.stderr if e.stderr else str(e)}"}
        except Exception as e:
            logger.error(f"处理grep请求时出错: {str(e)}")
            logger.error(traceback.format_exc())
            return {"error": f"处理请求时出错: {str(e)}"}
    
    async def handle_convert(params: Dict[str, Any]) -> Dict[str, Any]:
        """处理序列格式转换请求"""
        logger.info(f"处理convert请求: {params}")
        
        if "input_file" not in params:
            return {"error": "需要提供输入文件路径(input_file)"}
        if "output_format" not in params:
            return {"error": "需要提供输出格式(output_format)，可选：fasta, fastq"}
        
        input_file = params["input_file"]
        output_format = params["output_format"].lower()
        
        if output_format not in ["fasta", "fastq"]:
            return {"error": "输出格式必须是 fasta 或 fastq"}
        
        args = [SEQKIT_PATH, "convert"]
        
        if output_format == "fasta":
            args.append("-f" if params.get("full_header", True) else "-FF")
        elif output_format == "fastq":
            args.append("-q")
        
        output_file = params.get("output_file")
        if output_file:
            args.extend(["-o", output_file])
        
        try:
            if output_file:
                result = subprocess.run(
                    args + [input_file],
                    capture_output=True,
                    text=True,
                    check=True
                )
                return {"result": f"已保存到 {output_file}", "message": result.stdout}
            else:
                with tempfile.NamedTemporaryFile(mode="w+", delete=False) as temp:
                    temp_name = temp.name
                
                subprocess.run(
                    args + ["-o", temp_name, input_file],
                    check=True
                )
                
                with open(temp_name, "r") as f:
                    content = f.read()
                
                os.unlink(temp_name)
                return {"result": content}
        except subprocess.CalledProcessError as e:
            logger.error(f"执行convert命令时出错: {str(e)}")
            return {"error": f"执行错误: {e.stderr if e.stderr else str(e)}"}
        except Exception as e:
            logger.error(f"处理convert请求时出错: {str(e)}")
            logger.error(traceback.format_exc())
            return {"error": f"处理请求时出错: {str(e)}"}
    
    async def handle_translate(params: Dict[str, Any]) -> Dict[str, Any]:
        """处理核酸翻译请求"""
        logger.info(f"处理translate请求: {params}")
        
        if "input_file" not in params:
            return {"error": "需要提供输入文件路径(input_file)"}
        
        input_file = params["input_file"]
        args = [SEQKIT_PATH, "translate"]
        
        if params.get("frame"):
            args.extend(["--frame", str(params["frame"])])
        if params.get("trim"):
            args.append("--trim")
        if params.get("clean"):
            args.append("--clean")
        
        output_file = params.get("output_file")
        if output_file:
            args.extend(["-o", output_file])
        
        try:
            if output_file:
                result = subprocess.run(
                    args + [input_file],
                    capture_output=True,
                    text=True,
                    check=True
                )
                return {"result": f"已保存到 {output_file}", "message": result.stdout}
            else:
                with tempfile.NamedTemporaryFile(mode="w+", delete=False) as temp:
                    temp_name = temp.name
                
                subprocess.run(
                    args + ["-o", temp_name, input_file],
                    check=True
                )
                
                with open(temp_name, "r") as f:
                    content = f.read()
                
                os.unlink(temp_name)
                return {"result": content}
        except subprocess.CalledProcessError as e:
            logger.error(f"执行translate命令时出错: {str(e)}")
            return {"error": f"执行错误: {e.stderr if e.stderr else str(e)}"}
        except Exception as e:
            logger.error(f"处理translate请求时出错: {str(e)}")
            logger.error(traceback.format_exc())
            return {"error": f"处理请求时出错: {str(e)}"}

    async def main():
        logger.info("启动Seqkit MCP服务器")
        if not os.path.exists(SEQKIT_PATH):
            logger.error(f"错误: seqkit 可执行文件在 {SEQKIT_PATH} 不存在")
            print(f"错误: seqkit 可执行文件在 {SEQKIT_PATH} 不存在", file=sys.stderr)
            print("请从 https://github.com/shenwei356/seqkit/releases 下载适合您系统的版本", file=sys.stderr)
            print("并将其放置在与此脚本相同的目录中", file=sys.stderr)
            sys.exit(1)
        
        logger.info("开始设置stdio服务器")
        try:
            async with stdio_server() as streams:
                logger.info("获取到streams，准备运行app")
                await app.run(
                    streams[0],
                    streams[1],
                    InitializationOptions(
                        server_name="seqkit",
                        server_version="1.0.0",
                        capabilities=app.get_capabilities(
                            notification_options=NotificationOptions(),
                            experimental_capabilities={},
                        ),
                    )
                )
                logger.info("app.run完成")
        except Exception as e:
            logger.error(f"服务器运行时出错: {str(e)}")
            logger.error(traceback.format_exc())
            raise

    if __name__ == "__main__":
        try:
            logger.info("启动主程序")
            asyncio.run(main())
            logger.info("主程序正常结束")
        except KeyboardInterrupt:
            logger.info("用户中断程序")
            print("程序被用户中断")
        except Exception as e:
            logger.error(f"主程序运行时出错: {str(e)}")
            logger.error(traceback.format_exc())
            print(f"程序出错: {str(e)}")
except Exception as e:
    logger.critical(f"全局异常: {str(e)}")
    logger.critical(traceback.format_exc())
    print(f"全局异常: {str(e)}")
    print(traceback.format_exc()) 